from contextlib import asynccontextmanager
from functools import wraps

from sqlalchemy.ext.asyncio import AsyncSession

from server.config.database import AsyncSessionLocal
from server.config.secondary_databse import SecondaryAsyncSessionLocal


@asynccontextmanager
async def async_session_scope() -> AsyncSession:
    """Provide an asynchronous transactional scope around a series of operations."""
    session = AsyncSessionLocal()
    try:
        yield session
        await session.commit()
    except Exception:
        await session.rollback()
        raise
    finally:
        await session.close()


@asynccontextmanager
async def secondary_async_session_scope() -> AsyncSession:
    """Provide an asynchronous transactional scope around a series of operations."""
    session = SecondaryAsyncSessionLocal()
    try:
        yield session
        await session.commit()
    except Exception:
        await session.rollback()
        raise
    finally:
        await session.close()


def with_session(session_key='default'):
    """
    A decorator that provides a database session to an async function and manages transactions.

    :param session_key: A key to determine which session scope to use. Options are 'default' or 'secondary'.
    """

    def decorator(f):
        @wraps(f)
        async def wrapper(*args, **kwargs):
            # Choose the correct session scope based on session_key
            if session_key == 'secondary':
                session_scope = secondary_async_session_scope
            else:
                session_scope = async_session_scope

            async with session_scope() as session:
                try:
                    result = await f(session=session, *args, **kwargs)  # 使用 await 调用异步函数
                    await session.commit()  # 异步提交
                    return result
                except Exception as e:  # 捕获所有异常并回滚
                    await session.rollback()  # 异步回滚
                    raise e

        return wrapper

    return decorator
